import rclpy
from rclpy.node import Node
from nav_msgs.msg import Path
import yaml
import math
import sys

class PathStealer(Node):
    def __init__(self):
        super().__init__('path_stealer')
        
        # 在这里修改你想要的保存路径和文件名
        self.filename = '/home/jingzedong/B2_ws/src/b2_nav_bringup/params/tunnel_final_path1.yaml'

        # 订阅 Nav2 发布给 RViz 看的路径话题
        # 通常是 /plan，如果你用了平滑器可能是 /smoothed_plan
        self.subscription = self.create_subscription(
            Path,
            '/plan',
            self.listener_callback,
            1)
        
        self.interval = 3.5  # 每隔 2 米保存一个点
        self.get_logger().info('【就绪】正在监听 /plan 话题...')
        self.get_logger().info(f'路径将保存到: {self.filename}')
        self.get_logger().info('请在 RViz 中点击 "2D Nav Goal" 设置终点。')
        self.get_logger().info('一旦路径在 RViz 中出现，脚本将自动捕获并保存。')

    def listener_callback(self, msg):
        if not msg.poses:
            return

        self.get_logger().info(f'捕获到路径！原始点数: {len(msg.poses)}')
        
        # 处理并保存
        self.process_and_save(msg)
        
        # 保存完直接退出，防止重复保存
        self.get_logger().info('任务完成，程序退出。')
        raise SystemExit

    def process_and_save(self, path_msg):
        final_waypoints = []
        raw_poses = path_msg.poses
        
        # 加入起点
        last_added_pose = raw_poses[0]
        final_waypoints.append(self.pose_to_dict(last_added_pose))
        
        accumulated_dist = 0.0

        for i in range(1, len(raw_poses)):
            curr_pose = raw_poses[i]
            prev_pose = raw_poses[i-1]

            # 计算欧氏距离
            dist = math.sqrt(
                (curr_pose.pose.position.x - prev_pose.pose.position.x)**2 +
                (curr_pose.pose.position.y - prev_pose.pose.position.y)**2
            )
            accumulated_dist += dist

            # 达到间隔距离，保存该点
            if accumulated_dist >= self.interval:
                # 计算朝向 (Yaw)
                dx = curr_pose.pose.position.x - last_added_pose.pose.position.x
                dy = curr_pose.pose.position.y - last_added_pose.pose.position.y
                yaw = math.atan2(dy, dx)
                
                # 转换为四元数
                cy = math.cos(yaw * 0.5)
                sy = math.sin(yaw * 0.5)
                
                # 修改当前点的朝向，使其指向下一个点
                curr_pose.pose.orientation.z = sy
                curr_pose.pose.orientation.w = cy
                curr_pose.pose.orientation.x = 0.0
                curr_pose.pose.orientation.y = 0.0
                
                final_waypoints.append(self.pose_to_dict(curr_pose))
                
                last_added_pose = curr_pose
                accumulated_dist = 0.0

        # 保存文件
        with open(self.filename, 'w') as f:
            yaml.dump({'waypoints': final_waypoints}, f)
            
        print(f"\n=============================================")
        print(f" 成功！已保存 {len(final_waypoints)} 个巡检点到文件：{self.filename}")
        print(f"=============================================\n")

    def pose_to_dict(self, pose_stamped):
        return {
            'x': float(pose_stamped.pose.position.x),
            'y': float(pose_stamped.pose.position.y),
            'z': float(pose_stamped.pose.orientation.z),
            'w': float(pose_stamped.pose.orientation.w)
        }

def main():
    rclpy.init()
    node = PathStealer()
    try:
        rclpy.spin(node)
    except SystemExit:
        pass
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()

if __name__ == '__main__':
    main()